package openai

import (
	"context"
	"fmt"
	gogpt "github.com/sashabaranov/go-openai"
	"github.com/spf13/viper"
)

type OpenAI struct {
	Token string
}

func NewOpenAI() *OpenAI {
	return &OpenAI{}
}

// GetToken 获取token
func (p *OpenAI) GetToken() string {
	p.Token = viper.GetString("openai.apiKey")
	return p.Token
}

// GetTextResult 获取文本结果
func (p *OpenAI) GetTextResult(prompt string) (string, error) {

	c := gogpt.NewClient(p.GetToken())
	fmt.Println(prompt)
	resp, err := c.CreateChatCompletion(
		context.Background(),
		gogpt.ChatCompletionRequest{
			Model: gogpt.GPT3Dot5Turbo,
			Messages: []gogpt.ChatCompletionMessage{
				{
					Role:    gogpt.ChatMessageRoleUser,
					Content: prompt,
				},
			},
		},
	)

	if err != nil {
		fmt.Printf("ChatCompletion error: %v\n", err)
		return err.Error(), nil
	}

	fmt.Println(resp.Choices[0].Message.Content)
	return resp.Choices[0].Message.Content, nil
}

// GetCodeResult 获取代码结果
func (p *OpenAI) GetCodeResult(prompt string) (string, error) {
	c := gogpt.NewClient(p.GetToken())
	ctx := context.Background()

	req := gogpt.CompletionRequest{
		Model:            gogpt.CodexCodeDavinci002,
		Temperature:      0,
		MaxTokens:        2048,
		TopP:             1,
		FrequencyPenalty: 0,
		PresencePenalty:  0,
		BestOf:           1,
		Prompt:           prompt,
	}
	resp, err := c.CreateCompletion(ctx, req)
	if err != nil {
		return "", err
	}
	fmt.Println(resp.Choices[0].Text)
	return resp.Choices[0].Text, nil
}
